#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>
#include <iostream>
#include <vector>

#include <unistd.h>
#include <sys/socket.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <errno.h>
#include <signal.h>

#include "include.h"
#include "utils.h"

using namespace std;


#define DEFAULT_PORT			5555
#define DEFAULT_ADDR			"192.168.1.1"
#define DEFAULT_REMOTE_PORT		5555
#define DEFAULT_CONNECT_TIMEOUT	5000	// 5sec
#define DEFAULT_SEND_TIMEOUT	1000	// 1sec
#define DEFAULT_RECV_TIMEOUT	5000	// 5sec	
#define RECV_BUFFER_SIZE		1024 * 32

#define WAIT_FOREVER		0xFFFF

typedef struct user_info_client
{
	user_info_client() : num_boxes(0) {}

	unsigned int num_boxes;
	std::vector<unsigned int> ids;
}user_info_client;

typedef struct box_info_client
{
	box_info_client() : num_files(0) {}

	unsigned int num_files;
	std::vector<unsigned int> ids;
}box_info_client;


class CFile
{
public:
	CFile(const char *filename)
	{
		m_fp = fopen(filename, "w");
		if (m_fp == NULL)
			throw std::exception();
	}
	~CFile()
	{
		if (m_fp != NULL)
		{
			fclose(m_fp);
			m_fp = NULL;
		}
	}

	int Write(const void *data, size_t size)
	{
		return fwrite(data, 1, size, m_fp); 
	}
private:
	FILE *m_fp;
};

int auth_process(int sock)
{
	char username[] = "alexey";
	char password[] = "bla-bla-bla";
	int status;

	box_trans_header hdr;
	memset(&hdr, 0, sizeof(hdr));
	hdr.cmd = CMD_AUTH;
	hdr.param1 = strlen(username) + 1;
	hdr.param2 = strlen(password) + 1;
	hdr.data_len = hdr.param1 + hdr.param2;
	hdr.time1970 = time(NULL);

	size_t len = sizeof(hdr) + hdr.data_len;
	unsigned char *buf = (unsigned char*)malloc(len);
	if (buf == NULL)
	{
		printf("ERROR: Unable to allocate %d bytes for auth buffer\n", len);
		return -1;
	}
	memcpy(buf, &hdr, sizeof(hdr));
	memcpy(buf + sizeof(hdr), username, hdr.param1);
	memcpy(buf + sizeof(hdr) + hdr.param1, password, hdr.param2);

	// send buffer
	int bytes = UtilsTCPSend(sock, (const void*)buf, len, DEFAULT_SEND_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
	if (bytes == -1)
	{
		printf("ERROR: Unable to send CMD_AUTH buffer, %d bytes\n", len);
		free(buf);
		return -1;
	}
	printf("Sent CMD_AUTH command, %d bytes\n", len);
	fflush(stdout);

	// receive answer
	memset(&hdr, 0, sizeof(hdr));
	bytes = UtilsTCPReceive(sock, (char*)&hdr, sizeof(hdr), DEFAULT_RECV_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
	if (bytes == -1)
	{
		if (status == UTILS_TIMEOUT)
			printf("ERROR: UtilsTCPReceive timeout, %d bytes\n", sizeof(hdr));
		free(buf);
		return -1;
	}
	printf("Received CMD_AUTH answer, %d bytes\n", bytes);
	fflush(stdout);

	// check answer
	if (hdr.cmd == CMD_AUTH)
	{
		if (hdr.param1 == 0)	// auth success
		{
			printf("Auth success\n");
			free(buf);
			return 0;
		}
		else
		{
			printf("Auth failed\n");
			free(buf);
			return -1;
		}
	}
	else
	{
		printf("Incorrect hdr.cmd field\n");
		free(buf);
		return -1;
	}
}

// Get user info
int get_user_info(int sock, user_info_client& uinfo)
{
	int bytes = 0;
	int status;

	// Build header
	box_trans_header hdr;
	memset(&hdr, 0, sizeof(hdr));
	hdr.cmd = CMD_GET_USER_INFO;
	hdr.time1970 = time(NULL);

	// send command
	bytes = UtilsTCPSend(sock, (const void*)&hdr, sizeof(hdr), DEFAULT_SEND_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
	if (bytes == -1)
	{
		printf("ERROR: Unable to send CMD_GET_USER_INFO command\n");
		return -1;
	}
	printf("Sent CMD_GET_USER_INFO command, %d bytes\n", bytes);
	fflush(stdout);

	// receive answer
	bytes = UtilsTCPReceive(sock, (char*)&hdr, sizeof(hdr), DEFAULT_RECV_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
	if (bytes == -1)
	{
		if (status == UTILS_TIMEOUT)
			printf("ERROR: UtilsTCPReceive timeout, %d bytes\n", sizeof(hdr));
		return -1;
	}
	printf("Received CMD_GET_USER_INFO header\n");
	printf("	cmd - [%lu]\n", hdr.cmd);
	printf("	param1 - [%lu]\n", hdr.param1);
	printf("	param2 - [%lu]\n", hdr.param2);
	printf("	data_len - [%lu]\n", hdr.data_len);
	printf("	time1970 - [%lu]\n", hdr.time1970);
	fflush(stdout);

	if (hdr.cmd == CMD_GET_USER_INFO && hdr.data_len > 0)
	{
		char *buf = (char*)malloc(hdr.data_len);
		if (buf == NULL)
		{
			printf("ERROR: Unable to allocate %d bytes for buffer\n", hdr.data_len);
			return -1;
		}

		bytes = UtilsTCPReceive(sock, buf, hdr.data_len, DEFAULT_RECV_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
		if (bytes == -1)
		{
			if (status == UTILS_TIMEOUT)
				printf("ERROR: UtilsTCPReceive timeout, %d bytes\n", hdr.data_len);
			free(buf);
			return -1;
		}

		printf("Received CMD_GET_USER_INFO data\n");
		for (int i = 0; i < hdr.data_len; ++i)
			printf("%02X ", *(buf + i));
		printf("\n");

		unsigned int *data = (unsigned int*)buf;
		uinfo.num_boxes = *data;
		data++;
		for (int i = 0; i < uinfo.num_boxes; ++i)
		{
			uinfo.ids.push_back(*data);
			data++;
		}

		free(buf);
	}
	else
	{
		printf("ERROR: incorrect data_len\n");
		return -1;
	}

	return 0;
}

// Get box info
int get_box_info(int sock, unsigned int id, box_info_client& binfo)
{
	int bytes = 0;
	int status;

	// Build header
	box_trans_header hdr;
	memset(&hdr, 0, sizeof(hdr));
	hdr.cmd = CMD_GET_BOX_INFO;
	hdr.param1 = id;
	hdr.time1970 = time(NULL);

	// send command
	bytes = UtilsTCPSend(sock, (const void*)&hdr, sizeof(hdr), DEFAULT_SEND_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
	if (bytes == -1)
	{
		printf("ERROR: Unable to send CMD_GET_BOX_INFO command\n");
		return -1;
	}
	printf("Sent CMD_GET_BOX_INFO command, %d bytes\n", bytes);
	fflush(stdout);

	// receive answer
	bytes = UtilsTCPReceive(sock, (char*)&hdr, sizeof(hdr), DEFAULT_RECV_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
	if (bytes == -1)
	{
		if (status == UTILS_TIMEOUT)
			printf("ERROR: UtilsTCPReceive timeout, %d bytes\n", sizeof(hdr));
		return -1;
	}
	printf("Received CMD_GET_BOX_INFO header");
	printf("	cmd - [%lu]\n", hdr.cmd);
	printf("	param1 - [%lu]\n", hdr.param1);
	printf("	param2 - [%lu]\n", hdr.param2);
	printf("	data_len - [%lu]\n", hdr.data_len);
	printf("	time1970 - [%lu]\n", hdr.time1970);
	fflush(stdout);

	if (hdr.cmd == CMD_GET_BOX_INFO && hdr.data_len > 0)
	{
		char *buf = (char*)malloc(hdr.data_len);
		if (buf == NULL)
		{
			printf("ERROR: Unable to allocate %d bytes for buffer\n", hdr.data_len);
			return -1;
		}

		bytes = UtilsTCPReceive(sock, buf, hdr.data_len, DEFAULT_RECV_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
		if (bytes == -1)
		{
			if (status == UTILS_TIMEOUT)
				printf("ERROR: UtilsTCPReceive timeout, %d bytes\n", hdr.data_len);
			free(buf);
			return -1;
		}

		printf("Received CMD_GET_BOX_INFO data");
		for (int i = 0; i < hdr.data_len; ++i)
			printf("%02X ", *(buf + i));
		printf("\n");

		unsigned int *data = (unsigned int*)buf;
		binfo.num_files = *data;
		data++;
		for (int i = 0; i < binfo.num_files; ++i)
		{
			binfo.ids.push_back(*data);
			data++;
		}
		free(buf);
	}
	else
	{
		printf("ERROR: incorrect data_len\n");
		return -1;
	}

	return 0;
}

// Get file data
int get_file_data(int sock, unsigned int id, char *outfile)
{
	int bytes;
	unsigned long data_len;
	size_t len;
	int status;

	try
	{
		CFile file(outfile);

		box_trans_header hdr;
		memset(&hdr, 0, sizeof(hdr));
		hdr.cmd = CMD_READ_FILE_DATA;
		hdr.param1 = id;
		hdr.time1970 = time(NULL);

		// send command
		bytes = UtilsTCPSend(sock, (const void*)&hdr, sizeof(hdr), DEFAULT_SEND_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
		if (bytes == -1)
		{
			printf("ERROR: Unable to send CMD_READ_FILE_DATA command\n");
			return -1;
		}
		printf("Sent CMD_READ_FILE_DATA command, %d bytes\n", bytes);
		fflush(stdout);


		// Receive header answer
		bytes = UtilsTCPReceive(sock, (char*)&hdr, sizeof(hdr), DEFAULT_RECV_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
		if (bytes == -1)
		{
			if (status == UTILS_TIMEOUT)
				printf("ERROR: UtilsTCPReceive timeout, %d bytes\n", sizeof(hdr));
			return -1;
		}
		printf("Received header, cmd - %lu, param1 - %lu, param2 - %lu, data_len - %lu\n",
				hdr.cmd, hdr.param1, hdr.param2, hdr.data_len);
		fflush(stdout);


		if (hdr.cmd == CMD_START_FILE_DATA)
		{
			// allocate buffer for receiving
			char *buf = (char*)malloc(RECV_BUFFER_SIZE);
			if (buf == NULL)
			{
				printf("Unable to allocate %d bytes for receive buffer\n", RECV_BUFFER_SIZE);
				return -1;
			}

			// Receive data packets
			data_len = hdr.data_len;
			while(data_len)
			{
				len = (hdr.data_len > RECV_BUFFER_SIZE) ? RECV_BUFFER_SIZE : hdr.data_len;
				bytes = UtilsTCPReceive(sock, buf, len, DEFAULT_RECV_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
				if (bytes == -1)
				{
					printf("%s\n", strerror(errno));
					if (status == UTILS_TIMEOUT)
						printf("ERROR: UtilsTCPReceive timeout, %d bytes\n", len);
					free(buf);
					return -1;
				}
				printf("Received data, %d bytes\n", bytes);
				fflush(stdout);

				// write to file
				if (file.Write(buf, len) == -1)
				{
					printf("ERROR: Failed to write to file\n");
					free(buf);
					return -1;
				}
				data_len -= len;
			}

			do 
			{
				// Receive header
				bytes = UtilsTCPReceive(sock, (char*)&hdr, sizeof(hdr), DEFAULT_RECV_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
				if (bytes == -1)
				{
					if (status == UTILS_TIMEOUT)
						printf("ERROR: UtilsTCPReceive timeout, %d bytes\n", sizeof(hdr));
					free(buf);
					return -1;
				}
				printf("Received header, cmd - %lu, param1 - %lu, param2 - %lu, data_len - %lu\n",
						hdr.cmd, hdr.param1, hdr.param2, hdr.data_len);
				fflush(stdout);

				if (hdr.cmd == CMD_FILE_DATA || hdr.cmd == CMD_END_FILE_DATA)
				{
					// receive file data
					data_len = hdr.data_len;
					while(data_len)
					{
						len = (hdr.data_len > RECV_BUFFER_SIZE) ? RECV_BUFFER_SIZE : hdr.data_len;
						//bytes = UtilsTCPReceive(sock, buf, len, DEFAULT_RECV_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL);
						bytes = tcp_cycle_receive(sock, buf, len, DEFAULT_RECV_TIMEOUT, &status);
						if (bytes == -1)
						{
							if (status == UTILS_TIMEOUT)
								printf("ERROR: UtilsTCPReceive timeout, %d bytes\n", len);
							free(buf);
							return -1;
						}
						printf("Received data, %d bytes\n", bytes);
						fflush(stdout);

						// write to file
						if (file.Write(buf, len) == -1)
						{
							printf("ERROR: Failed to write to file\n");
							free(buf);
							return -1;
						}
						data_len -= len;
					}
				}
			} while (hdr.cmd == CMD_FILE_DATA);

			free(buf);
		}
		else
		{
			printf("Incorrect hdr.cmd field in answer\n");
			return -1;
		}
	}
	catch(...)
	{
		printf("ERROR: Unable to open output file\n");
		return -1;
	}

	return 0;
}


int main(int argc, char* argv[])
{
	int sock = -1;
	int ret = 0;
	int status;

	if (argc < 3)
	{
		printf("Usage: box_client <IP address> <port>\n");
		return -1;
	}

	// Create socket
	sock = UtilsTCPSocket(atoi(argv[2]), &status, /*(void(*)(const char*,...))&printf*/NULL);
	if (sock == UTILS_ERROR)
	{
		printf("UtilsTCPSocket error, %s\n", strerror(errno));
		return -1;
	}

	// Connect to remote address
	if (UtilsSocketConnect(sock, argv[1], atoi(argv[2]), DEFAULT_CONNECT_TIMEOUT, &status, /*(void(*)(const char*,...))&printf*/NULL) != UTILS_OK)
	{
		printf("UtilsSocketConnect error, %s\n", strerror(errno));
		close(sock);
		return -1;
	}

	printf("Connected\n");
	fflush(stdout);

	// Auth process
	if ((ret = auth_process(sock)) != 0)
	{
		shutdown(sock, 2);
		close(sock);
		return ret;
	}

	int key = 0;
	int code;
	do
	{
		system("clear");
		cout << "1. CMD_GET_USER_INFO\n";
		cout << "2. CMD_GET_BOX_INFO\n";
		cout << "3. CMD_GET_FILE_DATA\n";
		cout << "4. Exit\n";

		code = getchar();
		getchar(); // for Enter after number
		if (code == '1')	// CMD_GET_USER_INFO
		{
			// Get user info
			user_info_client uinfo;
			if ((ret = get_user_info(sock, uinfo)) != 0)
				break;

			printf("User info:\n");
			printf("	num_boxes = [%d]\n", uinfo.num_boxes);
			std::vector<unsigned int>::iterator it = uinfo.ids.begin();
			for (; it != uinfo.ids.end(); ++it)
				printf("	id = [%d]\n", *it);
			fflush(stdout);

			printf("Press any key\n");
			getchar();
		}
		else if (code == '2')	// CMD_GET_BOX_INFO
		{
			unsigned int id;
			printf("num box: ");
			scanf("%d", &id);
			getchar(); // for Enter after number

			// Get box info
			box_info_client binfo;
			if ((ret = get_box_info(sock, id, binfo)) != 0)
				break;

			printf("Box info:\n");
			printf("	num_files = [%d]\n", binfo.num_files);
			std::vector<unsigned int>::iterator it = binfo.ids.begin();
			for (; it != binfo.ids.end(); ++it)
				printf("	id = [%d]\n", *it);
			fflush(stdout);

			printf("Press any key\n");
			getchar();
		}
		else if (code == '3')	// CMD_GET_FILE_DATA
		{
			unsigned int id;
			printf("num file: ");
			scanf("%d", &id);
			getchar(); // for Enter after number

			// Get file data
			char outfile[] = "./output";
			if ((ret = get_file_data(sock, id, outfile)) != 0)
				break;

			printf("Press any key\n");
			getchar();
		}
		else if (code == '4')
			break;
	}while (code != 13);

	shutdown(sock, 2);
	close(sock);
	return ret;
}

